import unittest
from warp_field import profile
import time

import torch.nn as nn
import torch
from warp_field.profile.mmdet3d import NVTXIterBasedRunner,NVTXRunnerHook
import logging

class DummyDataset(torch.utils.data.Dataset):
    def __len__(self): return 100
    def __getitem__(self, idx):
        time.sleep(0.1)
        return torch.randn(10), torch.randn(1)

class TestModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)
    
    def train_step(self, data_batch, optimizer):
        time.sleep(0.5)
        return dict(loss=0.5)

    def forward(self, x):
        return self.linear(x)

# 定义batch处理器
def batch_processor(model, data):
    inputs, labels = data
    outputs = model(inputs)
    loss = torch.nn.MSELoss()(outputs, labels)
    return {'loss': loss}

class MMDetRunnerTest(unittest.TestCase):
    def test_nvtx(self):
        train_loader = torch.utils.data.DataLoader(DummyDataset(), batch_size=4)

        model = TestModule()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

        runner = NVTXIterBasedRunner(
            model=model,
            # batch_processor=batch_processor,
            optimizer=optimizer,
            work_dir='./exp',
            logger=logging.Logger('test'),
            max_iters=10,
        )

        runner.run(
            data_loaders=[train_loader],
            workflow=[('train', 5)],
        )